{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dcd6670a",
   "metadata": {},
   "source": [
    "# Enabling Your GPU for a Solver in MUSICA\n",
    "\n",
    "This tutorial will show you to use utilize a GPU for your MUSICA work.</br>\n",
    "However, this tutorial will not cover how to efficiently use a GPU through parallelization; it will simply introduce getting a GPU set up to run your code.</br>\n",
    "<b>Note:</b> This tutorial requires you to have a Linux GPU-ready environment handy, such as a supercomputing node; it will fail otherwise."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e934c67",
   "metadata": {},
   "source": [
    "## 1. Creating a GPU Virtual Environment\n",
    "\n",
    "Running code on a GPU requires a different install protocol when setting up a virtual environment.</br>\n",
    "To do so, run these commands in your terminal:\n",
    "\n",
    "```\n",
    "conda create --name musica_gpu python=3.9\n",
    "conda activate musica_gpu\n",
    "pip install --upgrade setuptools pip wheel\n",
    "pip install nvidia-pyindex\n",
    "pip install musica[gpu]\n",
    "conda install ipykernel scikit-learn seaborn scipy dask\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd76107",
   "metadata": {},
   "source": [
    "## 2. Importing MUSICA\n",
    "\n",
    "Importing MUSICA is largerly the same, but with an additional <i>is_cuda_available()</i> function to verify that the GPU is running properly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fe92903",
   "metadata": {},
   "outputs": [],
   "source": [
    "import musica\n",
    "import musica.mechanism_configuration as mc\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import qmc\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "from musica.cuda import is_cuda_available"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27039537",
   "metadata": {},
   "source": [
    "As with creating the music_box environment in MusicBox's [Basic Workflow Tutorial](https://ncar.github.io/music-box/branch/main/tutorials/1.%20basic_workflow.html), this cell may be slow to run the first time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d82609a",
   "metadata": {},
   "source": [
    "## 3. Running a Basic Solver on GPU\n",
    "\n",
    "This code is a copy of the [Hypercube Tutorial](2.%20hypercube.ipynb), but with an if statement added outside the main code to verify that it is running on a GPU.</br>\n",
    "If you are seeing \"Error: No GPU Available\" being printed, that means a GPU was not detected; verify that your environment has a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "54dc9c1b",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'is_cuda_available' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mNameError\u001b[39m                                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[43mis_cuda_available\u001b[49m():\n\u001b[32m      2\u001b[39m     A = mc.Species(name=\u001b[33m\"\u001b[39m\u001b[33mA\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m      3\u001b[39m     B = mc.Species(name=\u001b[33m\"\u001b[39m\u001b[33mB\u001b[39m\u001b[33m\"\u001b[39m)\n",
      "\u001b[31mNameError\u001b[39m: name 'is_cuda_available' is not defined"
     ]
    }
   ],
   "source": [
    "if is_cuda_available():\n",
    "    A = mc.Species(name=\"A\")\n",
    "    B = mc.Species(name=\"B\")\n",
    "    C = mc.Species(name=\"C\")\n",
    "    species = [A, B, C]\n",
    "    gas = mc.Phase(name=\"gas\", species=species)\n",
    "\n",
    "    r1 = mc.Arrhenius(\n",
    "        name=\"A_to_B\",\n",
    "        A=4.0e-3,  # Pre-exponential factor\n",
    "        C=50,      # Activation energy (units assumed to be K)\n",
    "        reactants=[A],\n",
    "        products=[B],\n",
    "        gas_phase=gas\n",
    "    )\n",
    "\n",
    "    r2 = mc.Arrhenius(\n",
    "        name=\"B_to_C\",\n",
    "        A=4.0e-3,\n",
    "        C=50,\n",
    "        reactants=[B],\n",
    "        products=[C],\n",
    "        gas_phase=gas\n",
    "    )\n",
    "\n",
    "    mechanism = mc.Mechanism(\n",
    "        name=\"musica_micm_example\",\n",
    "        species=species,\n",
    "        phases=[gas],\n",
    "        reactions=[r1, r2]\n",
    "    )\n",
    "\n",
    "    solver = musica.MICM(mechanism = mechanism, solver_type = musica.SolverType.cuda_rosenbrock)\n",
    "\n",
    "    num_grid_cells = 100\n",
    "    state = solver.create_state(num_grid_cells)\n",
    "\n",
    "    ndim = 5\n",
    "    nsamples = num_grid_cells\n",
    "\n",
    "    # Create a Latin Hypercube sampler in the unit hypercube\n",
    "    sampler = qmc.LatinHypercube(d=ndim)\n",
    "\n",
    "    # Generate samples\n",
    "    sample = sampler.random(n=nsamples)\n",
    "\n",
    "    # Define bounds for each dimension\n",
    "    l_bounds = [275, 100753.3, 0, 0, 0] # Lower bounds\n",
    "    u_bounds = [325, 101753.3, 10, 10, 10] # Upper bounds\n",
    "\n",
    "    # Scale the samples to the defined bounds\n",
    "    sample_scaled = qmc.scale(sample, l_bounds, u_bounds)\n",
    "\n",
    "    temperatures = sample_scaled[:, 0]\n",
    "    pressures = sample_scaled[:, 1]\n",
    "    concentrations = {\n",
    "        \"A\": [],\n",
    "        \"B\": [],\n",
    "        \"C\": []\n",
    "    }\n",
    "    concentrations[\"A\"] = sample_scaled[:, 2]\n",
    "    concentrations[\"B\"] = sample_scaled[:, 3]\n",
    "    concentrations[\"C\"] = sample_scaled[:, 4]\n",
    "\n",
    "    state.set_conditions(temperatures, pressures)\n",
    "    state.set_concentrations(concentrations)\n",
    "    concentrations_solved = []\n",
    "    time_step_length = 1\n",
    "    sim_length = 60\n",
    "    curr_time = 0\n",
    "\n",
    "    while curr_time <= sim_length:\n",
    "        solver.solve(state, curr_time)\n",
    "        concentrations_solved.append(state.get_concentrations())\n",
    "        curr_time += time_step_length\n",
    "\n",
    "    def convert_results_all_cells():\n",
    "        concentrations_solved_pd = []\n",
    "        time = []\n",
    "        for i in range(0, sim_length + 1, time_step_length):\n",
    "            for j in range(0, num_grid_cells):\n",
    "                concentrations_solved_pd.append({species: concentration[j] for species, concentration in concentrations_solved[int(i/time_step_length)].items()})\n",
    "                time.append(i)\n",
    "        df = pd.DataFrame(concentrations_solved_pd)\n",
    "        df = df.rename(columns = {'A' : 'CONC.A.mol m-3', 'B' : 'CONC.B.mol m-3', 'C' : 'CONC.C.mol m-3'})\n",
    "        df['time.s'] = time\n",
    "        df['ENV.temperature.K'] = np.repeat(temperatures[0], (sim_length/time_step_length + 1.0) * num_grid_cells)\n",
    "        df['ENV.pressure.Pa'] = np.repeat(pressures[0], (sim_length/time_step_length + 1.0) * num_grid_cells)\n",
    "        df['ENV.air number density.mol m-3'] = np.repeat(state.get_conditions()['air_density'][0], (sim_length/time_step_length + 1.0) * num_grid_cells)\n",
    "        df = df[['time.s', 'ENV.temperature.K', 'ENV.pressure.Pa', 'ENV.air number density.mol m-3', 'CONC.A.mol m-3', 'CONC.B.mol m-3', 'CONC.C.mol m-3']]\n",
    "        return concentrations_solved_pd, df\n",
    "\n",
    "    concentrations_solved_pd, df = convert_results_all_cells()\n",
    "\n",
    "    sns.lineplot(data=df, x='time.s', y='CONC.A.mol m-3', errorbar=('ci', 95), err_kws={'alpha' : 0.4}, label='CONC.A.mol m-3')\n",
    "    sns.lineplot(data=df, x='time.s', y='CONC.B.mol m-3', errorbar=('ci', 95), err_kws={'alpha' : 0.4}, label='CONC.B.mol m-3')\n",
    "    sns.lineplot(data=df, x='time.s', y='CONC.C.mol m-3', errorbar=('ci', 95), err_kws={'alpha' : 0.4}, label='CONC.C.mol m-3')\n",
    "    plt.title('Average concentration with CI over time')\n",
    "    plt.ylabel('Concentration (mol m-3)')\n",
    "    plt.xlabel('Time (s)')\n",
    "    plt.legend(loc='center right')\n",
    "    plt.show()\n",
    "\n",
    "    min_y = []\n",
    "    max_y = []\n",
    "    for i in range(0, sim_length + 1, time_step_length):\n",
    "        min_y.append({species: np.min(concentration) for species, concentration in concentrations_solved[int(i/time_step_length)].items()})\n",
    "        max_y.append({species: np.max(concentration) for species, concentration in concentrations_solved[int(i/time_step_length)].items()})\n",
    "    time_x = list(map(float, range(0, sim_length + 1, time_step_length)))\n",
    "\n",
    "    plt.fill_between(time_x, [y['A'] for y in min_y], [y['A'] for y in max_y], alpha = 0.4, label='CONC.A.mol m-3')\n",
    "    plt.fill_between(time_x, [y['B'] for y in min_y], [y['B'] for y in max_y], alpha = 0.4, label='CONC.B.mol m-3')\n",
    "    plt.fill_between(time_x, [y['C'] for y in min_y], [y['C'] for y in max_y], alpha = 0.4, label='CONC.C.mol m-3')\n",
    "    plt.title('Concentration range over time')\n",
    "    plt.ylabel('Concentration (mol m-3)')\n",
    "    plt.xlabel('Time (s)')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "else:\n",
    "    print(\"Error: No GPU Available\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "musica",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}